{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "# Kinematic Trajectory Optimization\n",
    "\n",
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/trajectories.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    BsplineTrajectory,\n",
    "    DiagramBuilder,\n",
    "    KinematicTrajectoryOptimization,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    MinimumDistanceLowerBoundConstraint,\n",
    "    Parser,\n",
    "    PositionConstraint,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    Role,\n",
    "    Solve,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import PublishPositionTrajectory\n",
    "from manipulation.scenarios import AddIiwa, AddPlanarIiwa, AddShape, AddWsg\n",
    "from manipulation.utils import ConfigureParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reaching into the shelves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trajopt_shelves_demo(avoid_collisions=True):\n",
    "    meshcat.Delete()\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddPlanarIiwa(plant)\n",
    "    wsg = AddWsg(plant, iiwa, roll=0.0, welded=True, sphere=True)\n",
    "    X_WStart = RigidTransform([0.8, 0, 0.65])\n",
    "    meshcat.SetObject(\"start\", Sphere(0.02), rgba=Rgba(0.9, 0.1, 0.1, 1))\n",
    "    meshcat.SetTransform(\"start\", X_WStart)\n",
    "    X_WGoal = RigidTransform([0.8, 0, 0.4])\n",
    "    meshcat.SetObject(\"goal\", Sphere(0.02), rgba=Rgba(0.1, 0.9, 0.1, 1))\n",
    "    meshcat.SetTransform(\"goal\", X_WGoal)\n",
    "\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    bin = parser.AddModelsFromUrl(\"package://manipulation/shelves.sdf\")[0]\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"shelves_body\", bin),\n",
    "        RigidTransform([0.88, 0, 0.4]),\n",
    "    )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(role=Role.kIllustration),\n",
    "    )\n",
    "    collision_visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(\n",
    "            prefix=\"collision\", role=Role.kProximity, visible_by_default=False\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    num_q = plant.num_positions()\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "\n",
    "    trajopt = KinematicTrajectoryOptimization(plant.num_positions(), 10)\n",
    "    prog = trajopt.get_mutable_prog()\n",
    "    trajopt.AddDurationCost(1.0)\n",
    "    trajopt.AddPathLengthCost(1.0)\n",
    "    trajopt.AddPositionBounds(\n",
    "        plant.GetPositionLowerLimits(), plant.GetPositionUpperLimits()\n",
    "    )\n",
    "    trajopt.AddVelocityBounds(\n",
    "        plant.GetVelocityLowerLimits(), plant.GetVelocityUpperLimits()\n",
    "    )\n",
    "\n",
    "    trajopt.AddDurationConstraint(0.5, 5)\n",
    "\n",
    "    # start constraint\n",
    "    start_constraint = PositionConstraint(\n",
    "        plant,\n",
    "        plant.world_frame(),\n",
    "        X_WStart.translation(),\n",
    "        X_WStart.translation(),\n",
    "        gripper_frame,\n",
    "        [0, 0.1, 0],\n",
    "        plant_context,\n",
    "    )\n",
    "    trajopt.AddPathPositionConstraint(start_constraint, 0)\n",
    "    prog.AddQuadraticErrorCost(\n",
    "        np.eye(num_q), q0, trajopt.control_points()[:, 0]\n",
    "    )\n",
    "\n",
    "    # goal constraint\n",
    "    goal_constraint = PositionConstraint(\n",
    "        plant,\n",
    "        plant.world_frame(),\n",
    "        X_WGoal.translation(),\n",
    "        X_WGoal.translation(),\n",
    "        gripper_frame,\n",
    "        [0, 0.1, 0],\n",
    "        plant_context,\n",
    "    )\n",
    "    trajopt.AddPathPositionConstraint(goal_constraint, 1)\n",
    "    prog.AddQuadraticErrorCost(\n",
    "        np.eye(num_q), q0, trajopt.control_points()[:, -1]\n",
    "    )\n",
    "\n",
    "    # start and end with zero velocity\n",
    "    trajopt.AddPathVelocityConstraint(\n",
    "        np.zeros((num_q, 1)), np.zeros((num_q, 1)), 0\n",
    "    )\n",
    "    trajopt.AddPathVelocityConstraint(\n",
    "        np.zeros((num_q, 1)), np.zeros((num_q, 1)), 1\n",
    "    )\n",
    "\n",
    "    # Solve once without the collisions and set that as the initial guess for\n",
    "    # the version with collisions.\n",
    "    result = Solve(prog)\n",
    "    if not result.is_success():\n",
    "        print(\"Trajectory optimization failed, even without collisions!\")\n",
    "        print(result.get_solver_id().name())\n",
    "    trajopt.SetInitialGuess(trajopt.ReconstructTrajectory(result))\n",
    "\n",
    "    if avoid_collisions:\n",
    "        # collision constraints\n",
    "        collision_constraint = MinimumDistanceLowerBoundConstraint(\n",
    "            plant, 0.01, plant_context, None, 0.1\n",
    "        )\n",
    "        evaluate_at_s = np.linspace(0, 1, 25)\n",
    "        for s in evaluate_at_s:\n",
    "            trajopt.AddPathPositionConstraint(collision_constraint, s)\n",
    "\n",
    "        def PlotPath(control_points):\n",
    "            traj = BsplineTrajectory(\n",
    "                trajopt.basis(), control_points.reshape((3, -1))\n",
    "            )\n",
    "            meshcat.SetLine(\n",
    "                \"positions_path\", traj.vector_values(np.linspace(0, 1, 50))\n",
    "            )\n",
    "\n",
    "        prog.AddVisualizationCallback(\n",
    "            PlotPath, trajopt.control_points().reshape((-1,))\n",
    "        )\n",
    "        result = Solve(prog)\n",
    "        if not result.is_success():\n",
    "            print(\"Trajectory optimization failed\")\n",
    "            print(result.get_solver_id().name())\n",
    "\n",
    "    PublishPositionTrajectory(\n",
    "        trajopt.ReconstructTrajectory(result), context, plant, visualizer\n",
    "    )\n",
    "    collision_visualizer.ForcedPublish(\n",
    "        collision_visualizer.GetMyContextFromRoot(context)\n",
    "    )\n",
    "\n",
    "\n",
    "trajopt_shelves_demo(avoid_collisions=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moving between bins (without hitting the cameras!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trajopt_bins_demo():\n",
    "    meshcat.Delete()\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    bin = parser.AddModelsFromUrl(\n",
    "        \"package://manipulation/two_bins_w_cameras.dmd.yaml\"\n",
    "    )[0]\n",
    "    iiwa = AddIiwa(plant, collision_model=\"with_box_collision\")\n",
    "    wsg = AddWsg(plant, iiwa, welded=True, sphere=True)\n",
    "    X_WStart = RigidTransform([0.5, 0, 0.15])\n",
    "    meshcat.SetObject(\"start\", Sphere(0.02), rgba=Rgba(0.9, 0.1, 0.1, 1))\n",
    "    meshcat.SetTransform(\"start\", X_WStart)\n",
    "    X_WGoal = RigidTransform([0, -0.6, 0.15])\n",
    "    meshcat.SetObject(\"goal\", Sphere(0.02), rgba=Rgba(0.1, 0.9, 0.1, 1))\n",
    "    meshcat.SetTransform(\"goal\", X_WGoal)\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(role=Role.kIllustration),\n",
    "    )\n",
    "    collision_visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(\n",
    "            prefix=\"collision\", role=Role.kProximity, visible_by_default=False\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    num_q = plant.num_positions()\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "\n",
    "    trajopt = KinematicTrajectoryOptimization(plant.num_positions(), 10)\n",
    "    prog = trajopt.get_mutable_prog()\n",
    "\n",
    "    q_guess = np.tile(q0.reshape((7, 1)), (1, trajopt.num_control_points()))\n",
    "    q_guess[0, :] = np.linspace(0, -np.pi / 2, trajopt.num_control_points())\n",
    "    path_guess = BsplineTrajectory(trajopt.basis(), q_guess)\n",
    "    trajopt.SetInitialGuess(path_guess)\n",
    "\n",
    "    # Uncomment this to see the initial guess:\n",
    "    PublishPositionTrajectory(path_guess, context, plant, visualizer)\n",
    "\n",
    "    trajopt.AddDurationCost(1.0)\n",
    "    trajopt.AddPathLengthCost(1.0)\n",
    "    trajopt.AddPositionBounds(\n",
    "        plant.GetPositionLowerLimits(), plant.GetPositionUpperLimits()\n",
    "    )\n",
    "    trajopt.AddVelocityBounds(\n",
    "        plant.GetVelocityLowerLimits(), plant.GetVelocityUpperLimits()\n",
    "    )\n",
    "\n",
    "    trajopt.AddDurationConstraint(0.5, 50)\n",
    "\n",
    "    # start constraint\n",
    "    start_constraint = PositionConstraint(\n",
    "        plant,\n",
    "        plant.world_frame(),\n",
    "        X_WStart.translation(),\n",
    "        X_WStart.translation(),\n",
    "        gripper_frame,\n",
    "        [0, 0.1, 0],\n",
    "        plant_context,\n",
    "    )\n",
    "    trajopt.AddPathPositionConstraint(start_constraint, 0)\n",
    "    prog.AddQuadraticErrorCost(\n",
    "        np.eye(num_q), q0, trajopt.control_points()[:, 0]\n",
    "    )\n",
    "\n",
    "    # goal constraint\n",
    "    goal_constraint = PositionConstraint(\n",
    "        plant,\n",
    "        plant.world_frame(),\n",
    "        X_WGoal.translation(),\n",
    "        X_WGoal.translation(),\n",
    "        gripper_frame,\n",
    "        [0, 0.1, 0],\n",
    "        plant_context,\n",
    "    )\n",
    "    trajopt.AddPathPositionConstraint(goal_constraint, 1)\n",
    "    prog.AddQuadraticErrorCost(\n",
    "        np.eye(num_q), q0, trajopt.control_points()[:, -1]\n",
    "    )\n",
    "\n",
    "    # start and end with zero velocity\n",
    "    trajopt.AddPathVelocityConstraint(\n",
    "        np.zeros((num_q, 1)), np.zeros((num_q, 1)), 0\n",
    "    )\n",
    "    trajopt.AddPathVelocityConstraint(\n",
    "        np.zeros((num_q, 1)), np.zeros((num_q, 1)), 1\n",
    "    )\n",
    "\n",
    "    # collision constraints\n",
    "    collision_constraint = MinimumDistanceLowerBoundConstraint(\n",
    "        plant, 0.001, plant_context, None, 0.01\n",
    "    )\n",
    "    evaluate_at_s = np.linspace(0, 1, 50)\n",
    "    for s in evaluate_at_s:\n",
    "        trajopt.AddPathPositionConstraint(collision_constraint, s)\n",
    "\n",
    "    result = Solve(prog)\n",
    "    if not result.is_success():\n",
    "        print(\"Trajectory optimization failed\")\n",
    "        print(result.get_solver_id().name())\n",
    "\n",
    "    PublishPositionTrajectory(\n",
    "        trajopt.ReconstructTrajectory(result), context, plant, visualizer\n",
    "    )\n",
    "    collision_visualizer.ForcedPublish(\n",
    "        collision_visualizer.GetMyContextFromRoot(context)\n",
    "    )\n",
    "\n",
    "\n",
    "trajopt_bins_demo()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
